--
-- Death knight rotation helper
--
if _G.Rotor.player_class ~= 'DEATHKNIGHT' then return end

-------------------------------------------
-- Local references to globals for speed and lint's happiness
-------------------------------------------
local Rotor = _G.Rotor

local UnitHealth = _G.UnitHealth
local UnitHealthMax = _G.UnitHealthMax
local strsplit = _G.strsplit
local max = _G.max
local min = _G.min
local GetTime = _G.GetTime
local GetRuneCooldown = _G.GetRuneCooldown
local GetRuneType = _G.GetRuneType
local UnitBuff = _G.UnitBuff
local UIParent = _G.UIParent
local CreateFrame = _G.CreateFrame
local select = _G.select

-- The module
local Mod = Rotor:NewModule('Death Knight')

-------------------------------------------
-- Constants
-------------------------------------------
local BLOOD_RUNE_TYPE   = 1
local UNHOLY_RUNE_TYPE  = 2
local FROST_RUNE_TYPE   = 3
local DEATH_RUNE_TYPE   = 4

local RUNE_COLORS = {
    {1,     0,      0}, -- Blood
    {0,     0.95,   0}, -- Unholy
    {0,     1,      1}, -- Frost
    {0.8,   0.1,    1}  -- Death
}

local SPEC_BLOOD = 1
local SPEC_FROST = 2
local SPEC_UNHOLY = 3

local BLOOD_PRESENCE_ID = 48263
local UNHOLY_PRESENCE_ID = 48265
local FROST_PRESENCE_ID = 48266

local FS_ID         = 49143 -- frost strike
local HB_ID         = 49184 -- howling blast
local IT_ID         = 45477 -- icy touch
local PS_ID         = 45462 -- plague strike
local OUTB_ID       = 77575 -- outbreak
local HOW_ID        = 57330 -- horn of winter
local OBLI_ID       = 49020 -- obliterate
local DS_ID         = 49998 -- death strike
local HS_ID         = 55050 -- heart strike
local RS_ID         = 56815 -- rune strike
local DND_ID        = 43265 -- death n decay
local MF_ID         = 47528 -- mind freeze
local AMS_ID        = 48707 -- anti magic shell
local FF_ID         = 59921 -- frost fever
local BP_ID         = 55078 -- blood plague
local RT_ID         = 48982 -- rune tap
local IBF_ID        = 48792 -- ice bound fortitude
local LB_ID         = 49039 -- lichborne
local DRW_ID        = 49028 -- dancing rune weapon
local DC_ID         = 47541 -- death coil
local BONESHIELD_ID = 49222 -- bone shield
local RD_ID         = 46584 -- rais dead
local DP_ID         = 48743 -- death pact
local FFOG_ID       = 59052 -- freezing fog
local KM_ID         = 51124 -- killng machine
local POF_ID        = 51271 -- pillar of frost
local AOD_ID        = 42650 -- army of dead
local BS_ID         = 77535 -- blood shield
local DG_ID         = 49576 -- death grip
local PEST_ID       = 50842 -- pestilence
local BB_ID         = 48721 -- blood boil
local CS_ID         = 81136 -- crimson scourge

local BLOOD_PRESENCE_NAME = Rotor.spell_name_cache[BLOOD_PRESENCE_ID]
local UNHOLY_PRESENCE_NAME = Rotor.spell_name_cache[UNHOLY_PRESENCE_ID]
local FROST_PRESENCE_NAME = Rotor.spell_name_cache[FROST_PRESENCE_ID]

local FS_NAME = Rotor.spell_name_cache[FS_ID]
local HB_NAME = Rotor.spell_name_cache[HB_ID]
local IT_NAME = Rotor.spell_name_cache[IT_ID]
local PS_NAME = Rotor.spell_name_cache[PS_ID]
local OUTB_NAME = Rotor.spell_name_cache[OUTB_ID]
local BS_NAME = Rotor.spell_name_cache[BS_ID]

-- spell costs in runic power
local fs_cost   = 40 -- frost strike
local drw_cost  = 60 -- dancing runic weapon
local dc_cost   = 40 -- death coil
local dp_cost   = 40 -- death pact
local rs_cost   = 30 -- runic strike
local mf_cost   = 20 -- mind freeze

local uhp_known = false -- unholy presence

--
-- Spell and aura defs
--
local spellAndAuraDefs = {
    common = {
        spells = {
            obli    = {OBLI_ID, rt1 = UNHOLY_RUNE_TYPE, rt2 = FROST_RUNE_TYPE}, -- Obliterate
            it      = {IT_ID, rt1 = FROST_RUNE_TYPE}, -- Icy touch
            ps      = {PS_ID, rt1 = UNHOLY_RUNE_TYPE}, -- Plague strike
            outb    = OUTB_ID, -- Outbreak
            how     = HOW_ID, -- Horn of winter
            ds      = {DS_ID, rt1 = UNHOLY_RUNE_TYPE, rt2 = FROST_RUNE_TYPE}, -- Death Strike
            rs      = RS_ID, -- Rune Strike
            dnd     = {DND_ID, rt1 = UNHOLY_RUNE_TYPE}, -- Death n Decay
            ibf     = {IBF_ID, icon_pos = 10, with_buff = true, off_gcd = true}, -- icebound fortitude
            mf      = {MF_ID, icon_pos = 11, off_gcd = true, range_check = true}, -- mind freeze
            ams     = {AMS_ID, icon_pos = 9, with_buff = true, off_gcd = true}, -- anti-magic shell
                      {PEST_ID, rt1 = BLOOD_RUNE_TYPE}, -- Pestilence
                      {BB_ID, rt1 = BLOOD_RUNE_TYPE}, -- Blood boil
        },
        auras = {
            how = {id = HOW_ID, from_others = true}, -- Horn of Winter
            ff = {id = FF_ID, is_debuff = true}, -- Frost fever
            bp = {id = BP_ID, is_debuff = true}, -- Blood plague
        },
    },
    blood = {
        spec = SPEC_BLOOD,
        spells = {
            boneshield  = {BONESHIELD_ID, rt1 = UNHOLY_RUNE_TYPE, icon_pos = 8}, -- Bone shield
            rt          = {RT_ID, rt1 = BLOOD_RUNE_TYPE, icon_pos = 7},    -- rune tap
            lb          = {LB_ID, icon_pos = 6, icon_type = 'reminder', min_power = 40},   -- lich borne
            drw         = {DRW_ID, icon_pos = 5, icon_type = 'reminder', min_power = 60, with_buff = true},   -- dancing rune weapon
            rd          = RD_ID, -- raise dead
            dp          = DP_ID, -- death pact
            hs          = {HS_ID, rt1 = BLOOD_RUNE_TYPE} -- Heart Strike
        },
        auras = {
            bs      = {id = BS_ID, icon_pos = 1}, -- blood shield
            cs      = CS_ID, -- crimson scourge
        },
    },
    frost = {
        spec = SPEC_FROST,
        spells = {
            hb =    {HB_ID, rt1 = FROST_RUNE_TYPE}, -- Howling blast
            fs =    FS_ID, -- Frost strike
                    {POF_ID, with_buff = true, icon_pos = 6}, -- pillar of frost
                    {RD_ID, icon_pos = 7}, -- raise dead
                    {AOD_ID, icon_pos = 5}, -- army of the dead (dont record consumed runes, channeled spell with long cast time)
                    {DG_ID, icon_pos = 8, range_check = true}, -- death grip
        },
        auras = {
            ffog = FFOG_ID, -- Freezing fog
            km = KM_ID, -- Killing machine
        },
    }
}

-- Local vars
local now = 0
local gcd = 0

local next_ff_check_time = 0
local next_bp_check_time = 0

local power_delta = 0

local strong_target = false

local rune_counts = {0, 0, 0, 0} -- array of ready runes count, by types (with count of gcd)
local consumed_runes = {0, 0, 0, 0} -- array of consumed runes, by types

local spells = {}
local buffs = {}
local debuffs = {}

local runes_text
----------------------------------------
function Mod:Init(opts)
    if not runes_text then
        runes_text = Rotor:createOverlayText('RotorRuneDisplay', 'BOTTOM')
        runes_text:SetFont(Rotor.cfg.mono_font, 13, 'THINOUTLINE')
        runes_text:SetWidth(Rotor.main_icon:GetWidth())
        runes_text:ClearAllPoints()
        runes_text:SetPoint('TOPLEFT', Rotor.main_icon, 'TOPLEFT', 0, 12)
    end
    runes_text:SetText('')
    runes_text:Show()

    Rotor.upper_text:SetFont(Rotor.cfg.condensed_font, Rotor.cfg.font_size, 'THINOUTLINE') -- shield display

    for i = 1, 4 do
        rune_counts[i] = 0
        consumed_runes[i] = 0
    end

    spells, buffs, debuffs = Rotor:RegisterSpellsAndAuras(spellAndAuraDefs)

    self.spells = spells
    self.buffs = buffs
    self.debuffs = debuffs

    dc_cost = Rotor:GetSpellCost(DC_ID)
    mf_cost = Rotor:GetSpellCost(MF_ID)
    if Rotor.spec == SPEC_BLOOD then
        drw_cost = Rotor:GetSpellCost(DRW_ID)
        dp_cost = Rotor:GetSpellCost(DP_ID)
        rs_cost = Rotor:GetSpellCost(RS_ID)
    elseif Rotor.spec == SPEC_FROST then
        fs_cost = Rotor:GetSpellCost(FS_ID)
    end

    if spells.mf then
        spells.mf.min_power = mf_cost
    end

    Rotor:initBlIndicator(4)
    if (Rotor.spec == SPEC_BLOOD) and (Rotor.player_level >= 80) then
        --[[
        bs_bar = Rotor:createBar('TOP', false, true, true)
        bs_bar:setColor({0.5, 0.5, 0})
        --]]
        Rotor:showUpperText('')
    end
    Rotor:showPowerDisplay()
    Rotor:setInterruptSpell(spells.mf)

    uhp_known = Rotor:HasSpell(UNHOLY_PRESENCE_ID)

    local spec_names = {'blood', 'frost', 'unholy'}
    Rotor:printOptInfo('Death knight module initialized ('..spec_names[Rotor.spec]..')')
    return true
end
----------------------------------------
function Mod:Done()
    spells = nil
    buffs = nil
    debuffs = nil
    self.spells = nil
    self.buffs = nil
    self.debuffs = nil
    runes_text:Hide()
end
----------------------------------------
local function HaveOneRune(rt)
    return (rune_counts[rt] - consumed_runes[rt] + rune_counts[DEATH_RUNE_TYPE] - consumed_runes[DEATH_RUNE_TYPE]) > 0
end
----------------------------------------
local function HaveTwoRunes(rt1, rt2)
    return (min(max(rune_counts[rt1] - consumed_runes[rt1], 0), 1) + min(max(rune_counts[rt2] - consumed_runes[rt2], 0), 1) + max(rune_counts[DEATH_RUNE_TYPE] - consumed_runes[DEATH_RUNE_TYPE], 0)) > 1
end
----------------------------------------
local function HaveTwoRunePairs(rt1, rt2)
    return (rune_counts[rt1]  - consumed_runes[rt1]) > 1 and (rune_counts[rt2] - consumed_runes[rt2]) > 1
end
----------------------------------------
local function HaveB()
    return HaveOneRune(BLOOD_RUNE_TYPE)
end
----------------------------------------
local function HaveF()
    return HaveOneRune(FROST_RUNE_TYPE)
end
----------------------------------------
local function HaveU()
    return HaveOneRune(UNHOLY_RUNE_TYPE)
end
----------------------------------------
local function HaveUF()
    return HaveTwoRunes(UNHOLY_RUNE_TYPE, FROST_RUNE_TYPE)
end
----------------------------------------
local last_runes_text = ""
local function rune_count_str(rt)
    local r, g, b = unpack(RUNE_COLORS[rt])
    return string.format("|cff%02x%02x%02x%s|r", r * 255, g * 255, b * 255, rune_counts[rt])
end
----------------------------------------
local function updateRunes()
    for i = 1, 4 do
        rune_counts[i] = 0
    end

    for i = 1, 6 do
        local start, cooldown = GetRuneCooldown(i)
        local cdtime = math.max(math.ceil(start + cooldown - now), 0)
        local rt = GetRuneType(i)
        if (start < now) and (cdtime - gcd <= 0) then
            rune_counts[rt] = rune_counts[rt] + 1
        end
    end

    local text = ""
    if Rotor.spec == SPEC_FROST then
        text = rune_count_str(DEATH_RUNE_TYPE)..rune_count_str(UNHOLY_RUNE_TYPE)..rune_count_str(FROST_RUNE_TYPE)
    else
        for i = 1, 4 do
            text = text..rune_count_str(i)
        end
    end

    if last_runes_text ~= text then
        runes_text:SetText(text)
        last_runes_text = text
    end
end
----------------------------------------
local function CheckUnsent(spell, time_since_last_sent)
    return spell and (spell.sentAt <= spell.doneAt) and (spell.doneAt + (time_since_last_sent or 0.5) < now)
end
----------------------------------------
local function CheckOffCdAndUnsent(spell, time_since_last_sent, wth)
    return spell and (spell.cd - (wth or 0.01) <= gcd) and CheckUnsent(spell, time_since_last_sent)
end
----------------------------------------
local function CheckOffCd(spell)
    return spell and (spell.cd <= gcd)
end
----------------------------------------
local function EnoughRunes(spell)
    if not spell.def.rt1 then
        return true
    elseif not spell.def.rt2 then
        return HaveOneRune(spell.def.rt1)
    else
        return HaveTwoRunes(spell.def.rt1, spell.def.rt2)
    end
end
----------------------------------------
local function ConsumeRunes(spell)
    if not spell or not spell.def.rt1 then
        return
    end
    local rt
    if rune_counts[spell.def.rt1] > 0 then
        rt = spell.def.rt1
    else
        rt = DEATH_RUNE_TYPE
    end
    consumed_runes[rt] = consumed_runes[rt] + 1
    if spell.def.rt2 then
        if rune_counts[spell.def.rt2] > 0 then
            rt = spell.def.rt2
        else
            rt = DEATH_RUNE_TYPE
        end
        consumed_runes[rt] = consumed_runes[rt] + 1
    end
end
----------------------------------------
local PLAGUE_CHECK_TIME_DELAY = 1
function Mod:CastSent(spell, spell_rec)
    if spell == FS_NAME then
        power_delta = fs_cost
    elseif spell == OUTB_NAME then
        next_ff_check_time = GetTime() + PLAGUE_CHECK_TIME_DELAY
        next_bp_check_time = next_ff_check_time
    elseif (spell == IT_NAME) or (spell == HB_NAME) then
        next_ff_check_time = GetTime() + PLAGUE_CHECK_TIME_DELAY
    elseif spell == PS_NAME then
        next_bp_check_time = GetTime() + PLAGUE_CHECK_TIME_DELAY
    end

    if spell_rec and spell_rec.def.rt1 then
        if ((spell_rec.id ~= IT_ID) and (spell_rec.id ~= HB_ID)) or not buffs.ffog or not buffs.ffog.active then -- todo: bloodboil and rune tap with buffs
            ConsumeRunes(spell_rec)
        end
    end
end
----------------------------------------
function Mod:CastFinished(spell, succ, fail_reason, spell_rec)
    if not succ then
        if spell == OUTB_NAME then
            next_ff_check_time = GetTime() + PLAGUE_CHECK_TIME_DELAY
            next_bp_check_time = next_ff_check_time
        elseif (spell == IT_NAME) or (spell == HB_NAME) then
            next_ff_check_time = GetTime() + PLAGUE_CHECK_TIME_DELAY
        elseif spell == PS_NAME then
            next_bp_check_time = GetTime() + PLAGUE_CHECK_TIME_DELAY
        end
    end

    if spell == FS_NAME then
        power_delta = 0
    end

    if spell_rec and spell_rec.def.rt1 then
        for i = 1, 4 do
            consumed_runes[i] = 0
        end
    end
end
----------------------------------------
local function CheckDiseases()
    local out_of_frost_runes = false

    if ((debuffs.ff.time_left < 1.5) and (now >= next_ff_check_time)) or
        (strong_target and (debuffs.bp.time_left < 1.5) and (now >= next_bp_check_time))
    then
        if strong_target and CheckOffCdAndUnsent(spells.outb, 0.5, 0.5) then
            return spells.outb.id
        end

        if (debuffs.ff.time_left < 1.5) and (now >= next_ff_check_time) then
            local spell = spells.hb or spells.it

            if CheckUnsent(spell, 10) then
                if HaveF() then
                    return spell.id
                else
                    out_of_frost_runes = true
                end
            end
        end

        if strong_target and (debuffs.bp.time_left < 1.5) and (now >= next_bp_check_time) and CheckUnsent(spells.ps, 10) and (HaveU() or out_of_frost_runes) then
            return spells.ps.id
        end
    end

    return 0
end
----------------------------------------
local function CheckPresence()
    if Rotor.spec == SPEC_BLOOD then
        if not UnitBuff('player', BLOOD_PRESENCE_NAME) then
            return BLOOD_PRESENCE_ID
        else
            return 0
        end
    elseif not uhp_known then
        if not UnitBuff('player', FROST_PRESENCE_NAME) then
            return FROST_PRESENCE_ID
        else
            return 0
        end
    elseif not UnitBuff('player', UNHOLY_PRESENCE_NAME) then
        return UNHOLY_PRESENCE_ID
    else
        return 0
    end
end
----------------------------------------
local function CheckHoW()
    return (buffs.how.time_left < 2) and CheckOffCdAndUnsent(spells.how, 3)
end
----------------------------------------
function Mod:NextSkill(gcd_arg, cur_spell)
    gcd = gcd_arg + 0.1
    now = GetTime()

    updateRunes()

    local spell_id

    -- Check presence
    spell_id = CheckPresence()
    if spell_id ~= 0 then
        return spell_id
    end

    if not Rotor.in_combat and CheckHoW() then
        return HOW_ID
    end

    if Rotor.spec == SPEC_FROST then
        spell_id = self:NextFrostSkill(gcd_arg, cur_spell)
    elseif Rotor.spec == SPEC_BLOOD then
        spell_id = self:NextBloodSkill(gcd_arg, cur_spell)
    else
        spell_id = self:NextUnholySkill(gcd_arg, cur_spell)
    end

    return spell_id
end
----------------------------------------
function Mod:NextFrostSkill(gcd_arg, cur_spell)
    -- Check plagues
    local spell_id = CheckDiseases()
    if spell_id ~= 0 then
        return spell_id
    end

    -- high prio obli: rune overcap or km proc
    if CheckUnsent(spells.obli) and
        (HaveTwoRunePairs(FROST_RUNE_TYPE, UNHOLY_RUNE_TYPE)
        or (rune_counts[DEATH_RUNE_TYPE] - consumed_runes[DEATH_RUNE_TYPE] > 2)
        or ((buffs.km.time_left > spells.obli.cd) and (spells.obli.cd - 1.5 < gcd_arg)))
    then
        return OBLI_ID
    end

    if (cur_spell ~= 0) then
        if (gcd_arg ~= 0) or self:CheckUsable(cur_spell) then
            return cur_spell
        end
    end

    -- fs overcap
    if strong_target and (Rotor.power - power_delta) >= 100 then
        return FS_ID
    end

    -- rime
    local hb_or_it = spells.hb or spells.it
    if (buffs.ffog.time_left > gcd_arg) and CheckUnsent(hb_or_it) then
        return hb_or_it.id
    end

    -- normal obli
    if HaveTwoRunes(FROST_RUNE_TYPE, UNHOLY_RUNE_TYPE) and CheckUnsent(spells.obli) then
        return OBLI_ID
    end

    -- frost strike
    if (Rotor.power - power_delta) >= fs_cost then
        return FS_ID
    end

    if CheckUnsent(hb_or_it) and HaveF() then
        return hb_or_it.id
    end

    if CheckOffCdAndUnsent(spells.how) then
        return HOW_ID
    end

    return 0
end
----------------------------------------
local function MinPowerForRuneStrike(hp)
    if CheckOffCd(spells.drw) then
        return drw_cost + rs_cost
    elseif CheckOffCd(spells.lb) and (hp <= 0.5) then
        return dc_cost + rs_cost
    else
        return mf_cost + rs_cost
    end
end
----------------------------------------
function Mod:NextBloodSkill(gcd_arg, cur_spell)
    -- death strike
    if HaveUF() and CheckUnsent(spells.ds) then
        return DS_ID
    end

    local hp = UnitHealth('player') / UnitHealthMax('player')

    if hp < 0.3 then
        -- OH SHIT
        if CheckOffCdAndUnsent(spells.rt) and spells.rt.usable then
            return RT_ID
        elseif CheckOffCdAndUnsent(spells.lb) and (Rotor.power >= dc_cost) then
            return LB_ID
        elseif CheckOffCdAndUnsent(spells.ibf) then
            return IBF_ID
        elseif Rotor.pet_alive and CheckOffCdAndUnsent(spells.dp) and (Rotor.power >= dp_cost) then
            return DP_ID
        elseif CheckOffCdAndUnsent(spells.rd) and (Rotor.power >= dp_cost) then
            return RD_ID
        elseif CheckOffCdAndUnsent(spells.boneshield) then
            return BONESHIELD_ID
        elseif CheckOffCdAndUnsent(spells.ams) then
            return AMS_ID
        end
        -- oh noooo, i gonna die!!
    end

    -- rune strike
    if (Rotor.power >= MinPowerForRuneStrike(hp)) then
        return RS_ID
    end

    -- outbreak
    if CheckOffCdAndUnsent(spells.outb) then
        return OUTB_ID
    end

    -- hearth strike
    local min_blood_runes
    if spells.rt.cd > 10 then
        min_blood_runes = 0
    else
        min_blood_runes = 1
    end
    if (rune_counts[BLOOD_RUNE_TYPE] > min_blood_runes) and CheckUnsent(spells.hs) then
        return HS_ID
    end

    if CheckOffCdAndUnsent(spells.drw) and (Rotor.power >= drw_cost) then
        return DRW_ID
    end

    if buffs.cs and buffs.cs.time_left > gcd_arg then
        return BB_ID
    end

    if CheckOffCdAndUnsent(spells.how) then
        return HOW_ID
    end

    return 0
end
----------------------------------------
function Mod:NextUnholySkill(gcd_arg, cur_spell)
    -- todo
    return 0
end
----------------------------------------
function Mod:SpellsAurasAndResourcesUpdated()
    if (Rotor.spec == SPEC_BLOOD) and (Rotor.player_level >= 80) then
        local bs_amount

        if buffs.bs.active then
            bs_amount = select(14, UnitBuff('player', BS_NAME)) or 0
        else
            bs_amount = 0
        end
        if bs_amount == 0 then
            Rotor:setUpperText('----')
        else
            if bs_amount >= 10000 then
                bs_amount = tostring(math.floor((bs_amount + 500) / 1000))..'k'
            else
                bs_amount = string.format('%.1fk', bs_amount / 1000)
            end
            Rotor:setUpperText(bs_amount)
        end
    end
end
----------------------------------------
function Mod:TargetChanged(old_target)
    -- reset disease sent time when switching targets
    spells.ps.sentAt = 0
    if spells.hb then
        spells.hb.sentAt = 0
    end
    spells.it.sentAt = 0

    if Rotor.on_target then
        runes_text:Show()
        strong_target = UnitHealth('target') > UnitHealth('player')*2
    else
        runes_text:Hide()
    end
end
----------------------------------------
function Mod:HandleCommandLine(cmd)
    local cmd, args = strsplit(' ', (cmd or ''):lower(), 2)
    -- handle command line here
end
----------------------------------------
function Mod:Detect()
    return Rotor.spec == SPEC_FROST or Rotor.spec == SPEC_BLOOD
end
